from aip import AipOcr
import os
import cv2
import numpy as np
from PIL import Image



# 将图片拆分保存
def png_file(png_path):
    file_name_list = os.listdir(png_path)
    file_num = 0
    for i in file_name_list :
        img = cv2.imread( "./images"+"/" + i )
        img_gray = cv2.cvtColor(img,cv2.COLOR_RGB2GRAY)  # 将图片灰度化
        screenshot_num = 0
        else_scerrnshot_num = 0
        while screenshot_num <= 38 and else_scerrnshot_num <=19:
            if file_num == 0:
                if screenshot_num <= 19 and screenshot_num != 6:
                    cv2.imwrite('scerrnshot/%s%s.png'%(file_num,screenshot_num), img_gray[22+22*screenshot_num:44+22*screenshot_num,15:162])
                    screenshot_num += 1
                elif screenshot_num > 19 and screenshot_num != 25:
                    cv2.imwrite('scerrnshot/%s%s.png'%(file_num,screenshot_num), img_gray[22+22*(screenshot_num-20):44+22*(screenshot_num-20),162:298])
                    screenshot_num += 1
                elif screenshot_num == 6 or screenshot_num == 25:
                    screenshot_num += 1
            else:
                if else_scerrnshot_num != 6:
                    cv2.imwrite('scerrnshot/%s%s.png'%(file_num,else_scerrnshot_num+19), img_gray[22+22*else_scerrnshot_num:44+22*else_scerrnshot_num,162:298])
                    else_scerrnshot_num += 1
                elif else_scerrnshot_num == 6:
                    else_scerrnshot_num += 1
        file_num += 1

# 将拆分后的图片二值化
def Bz_scerrnshot():
    scerrnshot_list = os.listdir("./scerrnshot")
    for i in scerrnshot_list :
        scerrnshot = cv2.imread("./scerrnshot/"+i)
        ret,thresh1 = cv2.threshold(scerrnshot,200,255,cv2.THRESH_BINARY)
        im = Image.fromarray(thresh1)
        im.save("./dealed_scerrnshot/%s"%i)

# 读取图片
def get_file_content(filePath):
    with open(filePath, 'rb') as fp:
        return fp.read()

# 逐个识别图片
def identify_file():
    dealed_file_list = os.listdir("./dealed_scerrnshot")
    dealed_file_list.sort(key=lambda x:int(x[:-4]))   # 给文件排序  以倒着数第四位'.'为分界线，按照‘.'左边的数字从小到大排序
    result_file = open("result.txt","a")
    write_num = 0
    for i in dealed_file_list :
        if write_num < 19:
            file_path = "./dealed_scerrnshot/" + i
            image = get_file_content(file_path)
            try:
                image_info =client.basicGeneral(image)
                while True:
                    if "error_code" in image_info :
                        if image_info["error_code"] ==  18 :
                            image_info =client.basicGeneral(image)
                        else:
                            break
                    else:
                        break
                if 'words_result' not in image_info:
                    result_file.write("NULL\t")
                elif len(image_info['words_result']) == 0:
                    result_file.write("NULL\t")
                elif 'words' not in image_info['words_result'][0]:
                    result_file.write("NULL\t")
                else:
                    word = image_info['words_result'][0]['words']
                    result_file.write("%s\t"%word)
            except Exception as err:
                print('[错误]ocr_image:', err)
                result_file.write("NULL\t")
            write_num += 1
        else :
            result_file.write("\n")
            print("读取成功一行")
            write_num = 0
    print("读取成功")


if __name__ == '__main__':
        # 我的 APPID AK SK #
    APP_ID = '22429528'
    API_KEY = 'WbeM5xWUBTevS5yClupzmk7M'
    SECRET_KEY = 'iz269aljZL0Lbj8b5ABNc5yxg4ZepeeC'
    png_path = "D:/hello Python/邓老师任务/ocr/images"
    client = AipOcr(APP_ID, API_KEY, SECRET_KEY)
    
    png_file(png_path)
    Bz_scerrnshot()
    identify_file()